import os
import time
import numpy as np
import pickle
import struct
import re
import cv2
from tqdm import tqdm
import imageio
import json
from PIL import Image
from multiprocessing import Pool
# import threadpool
# import threading
from collections import defaultdict
import random
import copy
from ctypes import CDLL,byref,create_string_buffer,cdll
from PIL import Image, ImageColor, ImageFont, ImageDraw, ImageFilter
import sys
import imgaug as ia
import imgaug.augmenters as iaa
import imgaug.parameters as iap
from imgaug.augmentables.polys import Polygon, PolygonsOnImage


from tool import filesystem,imgaug_tool,utils, data_enhance # export PYTHONPATH=$PYTHONPATH:`pwd`



PLATE_CHARS_PROVINCE = ["京", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑",
                        "苏", "浙", "皖", "闽", "赣", "鲁", "豫", "鄂", "湘", "粤",
                        "桂", "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁",
                        "新"]
# PLATE_CHARS_DIGIT = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
# PLATE_CHARS_LETTER = ["A", "B", "C", "D", "E", "F", "G",
#                         "H", "J", "K", "L", "M", "N",
#                         "P", "Q", "R", "S", "T",
#                         "U", "V", "W", "X", "Y", "Z"]

PLATE_CHARS_DIGIT = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"] \
                    + ["0"] * 3 \
                    + ["2"] * 3 \
                    + ["6"] * 3 \
                    + ["8"] * 3 \
                    + ["1"] * 3  

PLATE_CHARS_LETTER = ["A", "B", "C", "D", "E", "F", "G",
                        "H", "J", "K", "L", "M", "N",
                        "P", "Q", "R", "S", "T",
                        "U", "V", "W", "X", "Y", "Z"] \
                    + ["Q"] * 3 \
                    + ["Z"] * 3 \
                    + ["G"] * 3 \
                    + ["B"] * 3 \
                    + ["J"] * 2 \
                    + ["U"] * 2 \
                    + ["T"] * 2      


provinces = ["皖", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑", "苏", "浙", "京", "闽", "赣", "鲁", "豫", "鄂", "湘", "粤", "桂", "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁", "新", "警", "学", "O"]
alphabets = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'O']
ads = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'O']


# def init_sort_cpp():
#     # so_path = os.path.dirname(os.path.abspath(__file__)) + os.sep + "libsort_rect.so"     
#     so_path = "/home/swls/work_dir/git/rect_sort_cpp/build/src/libsort_rect.so"
#     rect_sort_so = cdll.LoadLibrary(so_path)
#     return rect_sort_so

# # rect_sort_so = init_sort_cpp()



# def get_point_cpp(points):

#     info_dict = []
#     for point in points:
#         i_d = {
#             "x":int(point[0]),
#             "y":int(point[1]),
#         }
#         info_dict.append(i_d)
#     input_info = json.dumps(info_dict)
#     # print(input_info)

#     so_path_buffer = create_string_buffer(input_info.encode('utf-8'), 65536*100)    
#     output_msg = create_string_buffer(65536*100)

#     out = rect_sort_so.get_point(byref(so_path_buffer),byref(output_msg))  
    
#         # print(output_msg.value.decode('utf-8'))
#     out_put_info = json.loads(output_msg.value.decode('utf-8'))

#     # return 0

#     new_points = [[out_put_info['x1'], out_put_info['y1']],
#                 [out_put_info['x2'], out_put_info['y2']],
#                 [out_put_info['x3'], out_put_info['y3']],
#                 [out_put_info['x4'], out_put_info['y4'] ]]

#     return new_points



def iou(rect1, rect2):
    y0 = np.max([rect1[1],rect2[1]])
    y1 = np.min([rect1[1] + rect1[3],rect2[1]+rect2[3]])

    x0 = np.max([rect1[0],rect2[0]])
    x1 = np.min([rect1[0] + rect1[2],rect2[0]+rect2[2]])
    return x1 > x0 and y1 > y0


def GenNewCh(f,val, w=45, h=90):
    img=Image.new("RGBA", (w,h),(0,0,0,0))
    draw = ImageDraw.Draw(img)
    color = (252,252,252) if PLATE_TYPE == 0 else (5,5,5)
    draw.text((0, 0),val,color,font=f)
    return img

class TextImageGenerator:
    def __init__(self):
        
        if PLATE_TYPE == 0 or PLATE_TYPE == 2:
            # # 128 测定值
            self.fontCE = ImageFont.truetype(fontChEng, 128, 0)    # 省简称使用字体
            # self.pil_img = Image.new("RGBA", (440, 140), (0, 0, 0,0))

        elif PLATE_TYPE == 1:
            # # 124 测定值
            self.fontCE = ImageFont.truetype(fontChEng, 124, 0)    # 省简称使用字体
            # self.pil_img = Image.new("RGBA", (480, 140), (0, 0, 0,0))
        else:
            raise NotImplementedError
        
        self.total_via_files = filesystem.get_all_filepath(dataset_dir, ["via_region_data_ori.json"])
        self.template_image_paths = [TEMPLATE_IMAGES+os.sep + f for f in os.listdir(TEMPLATE_IMAGES)]

    # 生成文本行图片
    def gen_text_line_image(self, pil_img, text, w=45, h=90):
        """
            # 蓝牌字体本身有上下间距
            # 绿牌字体没有上下左右间距
        """
        pts = []
        labels = []
        offset = 0
        # string_length = 8 if PLATE_TYPE == 1 else 7
        string_length = 7
        pil_img.paste(GenNewCh(self.fontCE, text[0], 45, 90), (offset+15, 25))
        pts.append([[offset+15, 25], [offset+15 + 45, 25], [offset+15 + 45, 25+90], [offset+15, 25+90]])
        labels.append(text[0])
        offset += 15+45
        for i in range(string_length-1):
            if i == 1:
                offset += 40 if PLATE_TYPE == 1 else 22
            offset += 9 if PLATE_TYPE == 1 else 12
            if text[i+1] != "x":
                w = 43 if PLATE_TYPE ==1 else 45
                pil_img.paste( GenNewCh(self.fontCE, text[i+1], w, 90), (offset, 25))
                pts.append([[offset, 25], [offset + w, 25], [offset+w, 25+90], [offset, 25+90]])
                labels.append(text[i+1])

            offset += 43 if PLATE_TYPE==1 else 45
        return pil_img, pts, labels


    def gen_plate_string(self):
        '''
        生成车牌号码字符串
        
        统一绿牌与蓝牌的字符长度，将绿牌后6位随机去掉一个
        '''
        plate_str = ""
        ## 增加重复出现的情况
        random_copy = np.random.randint(40) == 0

        c = np.random.choice(PLATE_CHARS_DIGIT) if np.random.randint(10) < 6 else np.random.choice(PLATE_CHARS_LETTER)

        # string_length = 8 if PLATE_TYPE == 1 else 7
        string_length = 7
        # 绿牌改成7位
        # random_remove_index = np.random.randint(2,8) if PLATE_TYPE ==1 else -1
        random_remove_index = -1

        for cpos in range(string_length):
            if cpos == 0:
                if ENCHANCE_京津冀:
                    indcies_list = list(np.arange(len(PLATE_CHARS_PROVINCE))) +[0]*5 + [2]*5 + [4]*5 
                    np.random.shuffle(indcies_list)
                    np.random.shuffle(indcies_list)
                    plate_str += PLATE_CHARS_PROVINCE[indcies_list[0]]
                else:
                    plate_str += np.random.choice(PLATE_CHARS_PROVINCE)

            elif cpos == 1:
                plate_str += np.random.choice(PLATE_CHARS_LETTER)
            else:
                # 绿牌 去掉一个数字
                if cpos == random_remove_index:
                    # 使用小写x 代替
                    plate_str += "x" 
                # 蓝牌
                else:
                    if random_copy and np.random.randint(5) != 0:
                        plate_str += c
                    elif np.random.randint(10) < 6:
                        plate_str += np.random.choice(PLATE_CHARS_DIGIT)
                    else :
                        plate_str += np.random.choice(PLATE_CHARS_LETTER)
        return plate_str

    def warp_perspective(self, img,max_x_ratio=0.55, max_y_ratio=0.4, max_h_ratio=0.2):
        """ 使图像轻微的畸变
            img 输入图像
            factor 畸变的参数
            size 为图片的目标尺寸
        """
        size_o = [img.shape[1],img.shape[0]]

        pts1 = np.float32([[0,0],[size_o[0],0],[size_o[0],size_o[1]],[0,size_o[1]]])
        
        # 车牌整体水平偏移, update ->  只会变短
        prob = list(np.arange(100)) + list(np.arange(60,100))  + list(np.arange(80,100))
        prob_choice  = np.random.choice(prob)
        # # 车牌右上 右下向左边收缩的偏移
        x_diff = int(float(size_o[0] * prob_choice)* 0.01 *max_x_ratio )
        # # 车牌右上 右下向上边抬的偏移
        y_diff = int(float(size_o[1] * prob_choice )* 0.01  *max_y_ratio )
        # # 车牌左上和右上点 左右摇摆的 偏移
        h_diff = int(float(size_o[1] * prob_choice )* 0.01  *max_h_ratio ) # int(size_o[0]* 0.1) # np.random.randint(int(img.shape[1]* 0.1))
        # 右下点往上抬的偏移
        v_diff = np.random.randint(int(size_o[1]* 0.1), int(size_o[1]* 0.25))

        # 去掉改左边的变换
        if False:
            # change left
            # x_diff = -x_diff if np.random.randint(2) < 1 else x_diff
            # y_diff = -y_diff if np.random.randint(2) < 1 else y_diff
            x_diff = x_diff
            y_diff = -y_diff
            # 水平偏移
            h_diff = np.random.randint(int(img.shape[1]* 0.05))
            if np.random.randint(2) < 1: h_diff = -h_diff

            p1 = [
                    0 ,
                    0 if y_diff < 0 else y_diff]
            p2 = [
                size_o[0]-x_diff+h_diff if x_diff < 0 else size_o[0]-x_diff+h_diff,
                  -y_diff  if y_diff < 0  else 0]
            p3 = [
                size_o[0]-x_diff if x_diff < 0 else size_o[0]-x_diff,
                size_o[1]-y_diff  if y_diff < 0  else size_o[1]
            ]
            # 竖直偏移
            v_diff = np.random.randint(int(img.shape[1]* 0.05))
            p4 =[
                0,
                0 + size_o[1]- v_diff if y_diff < 0 else y_diff + size_o[1] - v_diff
            ]
            pts2 = np.float32([p1,p2,p3,p4])
        else:
            # change right
            x_diff = -x_diff
            y_diff = -y_diff

            if np.random.randint(2) < 1: h_diff = -h_diff

            p1 = [
                    0 if h_diff < 0 else h_diff, 
                    -y_diff if y_diff < 0 else 0]
            p2 = [
                    size_o[0]+x_diff if h_diff < 0 else size_o[0]+x_diff+h_diff,
                    0  if y_diff < 0  else y_diff
                  ]
            p3 = [
                size_o[0]+x_diff-h_diff  if h_diff < 0 else size_o[0]+x_diff,
                  size_o[1] - v_diff if y_diff < 0  else size_o[1] +y_diff- v_diff 
            ]
            p4 =[
                -h_diff if h_diff < 0 else 0, 
                size_o[1] -y_diff if y_diff < 0 else  size_o[1]
            ]
            pts2 = np.float32([p1,p2,p3,p4])

        M  = cv2.getPerspectiveTransform(pts1,pts2)
        max_x ,max_y= np.max(pts2,axis=0 )

        dst = cv2.warpPerspective(img,M,(max_x, max_y))

        return dst, pts2

    def run(self, save_dir, start_index=0, save_count=300, images_in_per_dir=500, via_name="via_region_data_ori.json"):
        
        one_via_file = self.total_via_files[np.random.randint(len(self.total_via_files))]
        one_data_dict = filesystem.read_via_file(one_via_file)
        image_list = list(one_data_dict.keys())

        epochs = save_count // images_in_per_dir
        for e in range(epochs):
        
            new_save_dir = save_dir+ os.sep +str( start_index + e*images_in_per_dir)
            if os.path.exists(new_save_dir):
                print("continue: ",new_save_dir )
                continue
            else:
                os.makedirs(new_save_dir, exist_ok=True)
                
            save_info = []
            for i in tqdm(range(start_index +e*images_in_per_dir, start_index +(e+1)*images_in_per_dir)):
                
                background_path = image_list[i % len(image_list)]
                background_image = cv2.imread(background_path)
                de = data_enhance.CRNNPlateEnhance(background_image, data_dict=one_data_dict[background_path])

                for c in range(np.random.randint(2,6)):
                    # # 更新背景图
                    if PLATE_TYPE == 1:
                        pil_img = Image.new("RGBA", (428, 140), (0, 0, 0,0))
                    elif PLATE_TYPE == 0 or PLATE_TYPE == 2:
                        pil_img = Image.new("RGBA", (440, 140), (0, 0, 0,0))
                    else:
                        raise NotImplementedError
                    plate_str = self.gen_plate_string()
                    text_line_image, pts, labels = self.gen_text_line_image(pil_img, plate_str) 
                    # print(pts)
                    # return 0

                    background = Image.open(self.template_image_paths[np.random.randint(len(self.template_image_paths))]) # 060086
                    background.paste(text_line_image, (0,0), text_line_image)

                    ratio = np.random.choice(list(list(range(42,50)) + list(range(42,60)) + list(range(42,70)) )) * 0.01
                    text_image = background.resize((int(background.size[0] * ratio), int(background.size[1] *  ratio)), Image.ANTIALIAS)

                    new_image = text_image.convert('RGBA')
                    image_array = np.frombuffer(new_image.tobytes(), dtype=np.uint8)
                    image_array = image_array.reshape((new_image.size[1], new_image.size[0], 4))
                    image_array = image_array[:,:,:3]

                    cv_image = imgaug_tool.aug_image_for_plate(image_array)

                    save_path = new_save_dir + os.sep+ "{}_{}.jpg".format(i, c) 
                    de.random_paste_plate(cv_image, "".join(labels), save_path)

                save_info.append(de)

            # save
            
            save_list = []
            for idx, de in enumerate(save_info):
                
                sub_info = de.write_image()
                save_list.extend(sub_info)

            with open(new_save_dir + os.sep + "label.txt", "w") as wf:
                for line in save_list:
                    wf.write(os.path.basename(line[0]) + " "+ line[1] + "\n")


def thread_run(save_path, start_index, images_in_per_dir, save_count):
    tig = TextImageGenerator()
    # print("save_path: ", save_path)
    # print("start_index: ", start_index)

    tig.run(save_path,start_index=start_index, save_count=save_count,images_in_per_dir=images_in_per_dir)


if __name__ == "__main__":

    total_count = 600000
    save_count = 4000
    images_in_per_dir = 500
    thread_count = 5

    ENCHANCE_京津冀 = False
    # 0 -> 蓝牌  1 -> 绿牌  2 -> 黄牌
    PLATE_TYPE= int(sys.argv[1])

    save_dir = "/home/swls/work_dir/ocr/code/syn/syn_plate_300k"
    dataset_dir = "/home/swls/work_dir/ocr/code/yolo_train/yolo_plate_via"
    if PLATE_TYPE == 0:
        TEMPLATE_IMAGES = "/home/swls/tmp_data/plate_data/font/template"
        fontChEng = "/home/swls/tmp_data/plate_data/font/plate_blue.ttf"

    elif PLATE_TYPE == 1:
        TEMPLATE_IMAGES = "/home/swls/tmp_data/plate_data/font/new_green_template_2"
        fontChEng = "/home/swls/tmp_data/plate_data/font/vanlance.ttf"

    elif PLATE_TYPE == 2:
        TEMPLATE_IMAGES = "/home/swls/tmp_data/plate_data/font/yellow_template"
        fontChEng = "/home/swls/tmp_data/plate_data/font/plate_blue.ttf"
    
    else:
        pass

    start_index = int(sys.argv[2])
    # # start_index = int(0)
    thread_run(save_dir, start_index, images_in_per_dir, save_count)


    # # # # write to sh
    # file_dir = "/home/swls/work_dir/git/python_script/sunjie/plate_recognition"
    # filw_write = [open(file_dir+os.sep+"gen_crnn_plate_{}.sh".format(i), "w") for i in range(thread_count)]
        
    # for i in range(total_count // save_count  ):
    #     if i < total_count * 0.35 / save_count:
    #         plate_type = 0
    #     elif i < total_count * 0.70 / save_count:
    #         plate_type = 1
    #     else:
    #         plate_type = 2         
    #     filw_write[i % thread_count].write("python sunjie/plate_recognition/gen_crnn_plate.py {} {}".format(plate_type, i*save_count) + "\n")
    # for f in filw_write:
    #     f.close()

    # # print('Parent process %s.' % os.getpid())
    # p = Pool(thread_count)
    # for i in range(int(total_count / images_in_per_dir)):
    #     p.apply_async(thread_run, args=(save_dir, images_in_per_dir*i, images_in_per_dir ))
    # # print('Waiting for all subprocesses done...')
    # p.close()
    # p.join()
    # print('All subprocesses done.')




"""

cd /home/swls/work_dir/git/python_script && conda activate tf3 && export PYTHONPATH=$PYTHONPATH:`pwd` && sh sunjie/plate_recognition/gen_crnn_plate_0.sh

@2020-01-08     修改 warp_perspective 函数默认参数
                修改车牌背景图片颜色选取方式
                增强车牌相似字符
"""